#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/7/23
# @USER    : Shengji He
# @File    : dicom_io.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
import os
import zipfile
import platform
import shutil
import subprocess
import glob


def command_exists(command):
    return shutil.which(command) is not None


def download_dcm2niix(executable_path=None):
    import urllib.request
    print("  Downloading dcm2niix...")

    if platform.system() == "Windows":
        # url = "https://github.com/rordenlab/dcm2niix/releases/latest/download/dcm2niix_win.zip"
        url = "https://github.com/rordenlab/dcm2niix/releases/download/v1.0.20230411/dcm2niix_win.zip"
    elif platform.system() == "Darwin":  # Mac
        # raise ValueError("For MacOS automatic installation of dcm2niix not possible. Install it manually.")
        if platform.machine().startswith("arm") or platform.machine().startswith("aarch"):  # arm
            # url = "https://github.com/rordenlab/dcm2niix/releases/latest/download/macos_dcm2niix.pkg"
            url = "https://github.com/rordenlab/dcm2niix/releases/download/v1.0.20230411/dcm2niix_macos.zip"
        else:  # intel
            # unclear if this is the right link (is the same as for arm)
            # url = "https://github.com/rordenlab/dcm2niix/releases/latest/download/macos_dcm2niix.pkg"
            url = "https://github.com/rordenlab/dcm2niix/releases/download/v1.0.20230411/dcm2niix_macos.zip"
    elif platform.system() == "Linux":
        # url = "https://github.com/rordenlab/dcm2niix/releases/latest/download/dcm2niix_lnx.zip"
        url = "https://github.com/rordenlab/dcm2niix/releases/download/v1.0.20230411/dcm2niix_lnx.zip"
    else:
        raise ValueError("Unknown operating system. Can not download the right version of dcm2niix.")

    if executable_path is None:
        executable_path = get_weights_dir()

    zip_file = os.path.join(executable_path, "dcm2niix.zip")
    urllib.request.urlretrieve(url, zip_file)
    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
        zip_ref.extractall(executable_path)

    # Give execution permission to the script
    if platform.system() == "Windows":
        os.chmod(os.path.join(executable_path, "dcm2niix.exe"), 0o755)
    else:
        os.chmod(os.path.join(executable_path, "dcm2niix"), 0o755)

    # Clean up
    if os.path.exists(zip_file):
        os.remove(zip_file)
    if os.path.exists(os.path.join(executable_path, "dcm2niibatch")):
        os.remove(os.path.join(executable_path, "dcm2niibatch"))


def dcm_to_nifti_LEGACY(input_path, output_folder, out_name, verbose=False, executable_path=None, remove_json=True):
    """
    Uses dcm2niix (does not properly work on windows)

    input_path: a directory of dicom slices
    output_path: a nifti file path
    """
    if executable_path is None:
        executable_path = get_weights_dir()

    if command_exists("dcm2niix"):
        dcm2niix = "dcm2niix"
    else:
        if platform.system() == "Windows":
            dcm2niix = os.path.join(executable_path, "dcm2niix.exe")
        else:
            dcm2niix = os.path.join(executable_path, "dcm2niix")
        if not os.path.exists(dcm2niix):
            download_dcm2niix(executable_path)

    # subprocess.call(f"\"{dcm2niix}\" -o {output_folder} -z y -f {out_name} {input_path} {verbose}", shell=True)
    subprocess.call(f"\"{dcm2niix}\" -o \"{output_folder}\" -z i -f \"{out_name}\" -v {int(verbose)} \"{input_path}\"",
                    shell=True)

    # dcm2niix.exe -o D:\heshengji\iAC_DL_Data\Clinicdata_2019-2020\77407_CHU_LI_MEI_34729 -z i -f ac_body2 -v 1 "W:\clinic\PUMC\Clinicdata_2019-2020\Patient_CHU LI MEI_77407_20200814182337\CTDICOM"

    output_path = os.path.join(output_folder, out_name + '.nii.gz')

    if not os.path.exists(output_path):
        print(f"Content of dcm2niix output folder ({output_folder}):")
        print(glob.glob(os.path.join(output_folder, '*')))
        raise ValueError("dcm2niix failed to convert dicom to nifti.")

    nii_files = glob.glob(os.path.join(output_folder, '*.nii.gz'))

    if len(nii_files) > 1:
        print("WARNING: Dicom to nifti resulted in several nifti files. Skipping files which contain ROI in filename.")
        for nii_file in nii_files:
            # output file name is "converted_dcm.nii.gz" so if ROI in name, then this can be deleted
            if "ROI" in os.path.basename(nii_file):
                os.remove(nii_file)
                print(f"Skipped: {nii_file}")

    nii_files = glob.glob(os.path.join(output_folder, '*.nii.gz'))

    if len(nii_files) > 1:
        print("WARNING: Dicom to nifti resulted in several nifti files.")
        print([os.path.basename(f) for f in nii_files])
        while True:
            print("Only using first one? Yes or No (Y/N)")
            flag = input()
            if flag.lower().strip() in ['yes', 'y']:
                for nii_file in nii_files[1:]:
                    os.remove(nii_file)
                break
            elif flag.lower().strip() in ['no', 'n']:
                break
            else:
                pass
        # todo: have to rename first file to not contain any counter which is automatically added by dcm2niix

    if remove_json:
        os.remove(str(output_path)[:-7] + ".json")


def dcm_to_nifti(input_path, output_path, tmp_dir=None, verbose=False):
    """
    Uses dicom2nifti package (also works on windows)

    input_path: a directory of dicom slices or a zip file of dicom slices
    output_path: a nifti file path
    tmp_dir: extract zip file to this directory, else to the same directory as the zip file
    """
    # Check if input_path is a zip file and extract it
    import dicom2nifti

    if zipfile.is_zipfile(input_path):
        if verbose:
            print(f"Extracting zip file: {input_path}")
        extract_dir = os.path.splitext(input_path)[0] if tmp_dir is None else tmp_dir / "extracted_dcm"
        with zipfile.ZipFile(input_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
            input_path = extract_dir

    # Convert to nifti
    dicom2nifti.dicom_series_to_nifti(input_path, output_path, reorient_nifti=True)
